{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Deep Deterministic Policy Gradient (DDPG)\n",
    "\n",
    "**Author:** [amifunny](https://github.com/amifunny)<br>\n",
    "**Date created:** 2020/06/04<br>\n",
    "**Last modified:** 2020/06/06<br>\n",
    "**Description:** Implementing DDPG algorithm on the Inverted Pendulum Problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "**Deep Deterministic Policy Gradient (DDPG)** is a model-free off-policy algorithm for\n",
    "learning continous actions.\n",
    "\n",
    "It combines ideas from DPG (Deterministic Policy Gradient) and DQN (Deep Q-Network).\n",
    "It uses Experience Replay and slow-learning target networks from DQN, and it is based on DPG,\n",
    "which can operate over continuous action spaces.\n",
    "\n",
    "This tutorial closely follow this paper -\n",
    "[Continuous control with deep reinforcement learning](https://arxiv.org/pdf/1509.02971.pdf)\n",
    "\n",
    "## Problem\n",
    "\n",
    "We are trying to solve the classic **Inverted Pendulum** control problem.\n",
    "In this setting, we can take only two actions: swing left or swing right.\n",
    "\n",
    "What make this problem challenging for Q-Learning Algorithms is that actions\n",
    "are **continuous** instead of being **discrete**. That is, instead of using two\n",
    "discrete actions like `-1` or `+1`, we have to select from infinite actions\n",
    "ranging from `-2` to `+2`.\n",
    "\n",
    "## Quick theory\n",
    "\n",
    "Just like the Actor-Critic method, we have two networks:\n",
    "\n",
    "1. Actor - It proposes an action given a state.\n",
    "2. Critic - It predicts if the action is good (positive value) or bad (negative value)\n",
    "given a state and an action.\n",
    "\n",
    "DDPG uses two more techniques not present in the original DQN:\n",
    "\n",
    "**First, it uses two Target networks.**\n",
    "\n",
    "**Why?** Because it add stability to training. In short, we are learning from estimated\n",
    "targets and Target networks are updated slowly, hence keeping our estimated targets\n",
    "stable.\n",
    "\n",
    "Conceptually, this is like saying, \"I have an idea of how to play this well,\n",
    "I'm going to try it out for a bit until I find something better\",\n",
    "as opposed to saying \"I'm going to re-learn how to play this entire game after every move\".\n",
    "See this [StackOverflow answer](https://stackoverflow.com/a/54238556/13475679).\n",
    "\n",
    "**Second, it uses Experience Replay.**\n",
    "\n",
    "We store list of tuples `(state, action, reward, next_state)`, and instead of\n",
    "learning only from recent experience, we learn from sampling all of our experience\n",
    "accumulated so far.\n",
    "\n",
    "Now, let's see how is it implemented.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We use [OpenAIGym](http://gym.openai.com/docs) to create the environment.\n",
    "We will use the `upper_bound` parameter to scale our actions later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "problem = \"Pendulum-v0\"\n",
    "env = gym.make(problem)\n",
    "\n",
    "num_states = env.observation_space.shape[0]\n",
    "print(\"Size of State Space ->  {}\".format(num_states))\n",
    "num_actions = env.action_space.shape[0]\n",
    "print(\"Size of Action Space ->  {}\".format(num_actions))\n",
    "\n",
    "upper_bound = env.action_space.high[0]\n",
    "lower_bound = env.action_space.low[0]\n",
    "\n",
    "print(\"Max Value of Action ->  {}\".format(upper_bound))\n",
    "print(\"Min Value of Action ->  {}\".format(lower_bound))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "To implement better exploration by the Actor network, we use noisy perturbations, specifically\n",
    "an **Ornstein-Uhlenbeck process** for generating noise, as described in the paper.\n",
    "It samples noise from a correlated normal distribution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "class OUActionNoise:\n",
    "    def __init__(self, mean, std_deviation, theta=0.15, dt=1e-2, x_initial=None):\n",
    "        self.theta = theta\n",
    "        self.mean = mean\n",
    "        self.std_dev = std_deviation\n",
    "        self.dt = dt\n",
    "        self.x_initial = x_initial\n",
    "        self.reset()\n",
    "\n",
    "    def __call__(self):\n",
    "        # Formula taken from https://www.wikipedia.org/wiki/Ornstein-Uhlenbeck_process.\n",
    "        x = (\n",
    "            self.x_prev\n",
    "            + self.theta * (self.mean - self.x_prev) * self.dt\n",
    "            + self.std_dev * np.sqrt(self.dt) * np.random.normal(size=self.mean.shape)\n",
    "        )\n",
    "        # Store x into x_prev\n",
    "        # Makes next noise dependent on current one\n",
    "        self.x_prev = x\n",
    "        return x\n",
    "\n",
    "    def reset(self):\n",
    "        if self.x_initial is not None:\n",
    "            self.x_prev = self.x_initial\n",
    "        else:\n",
    "            self.x_prev = np.zeros_like(self.mean)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "The `Buffer` class implements Experience Replay.\n",
    "\n",
    "---\n",
    "![Algorithm](https://i.imgur.com/mS6iGyJ.jpg)\n",
    "---\n",
    "\n",
    "\n",
    "**Critic loss** - Mean Squared Error of `y - Q(s, a)`\n",
    "where `y` is the expected return as seen by the Target network,\n",
    "and `Q(s, a)` is action value predicted by the Critic network. `y` is a moving target\n",
    "that the critic model tries to achieve; we make this target\n",
    "stable by updating the Target model slowly.\n",
    "\n",
    "**Actor loss** - This is computed using the mean of the value given by the Critic network\n",
    "for the actions taken by the Actor network. We seek to maximize this quantity.\n",
    "\n",
    "Hence we update the Actor network so that it produces actions that get\n",
    "the maximum predicted value as seen by the Critic, for a given state.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "class Buffer:\n",
    "    def __init__(self, buffer_capacity=100000, batch_size=64):\n",
    "\n",
    "        # Number of \"experiences\" to store at max\n",
    "        self.buffer_capacity = buffer_capacity\n",
    "        # Num of tuples to train on.\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        # Its tells us num of times record() was called.\n",
    "        self.buffer_counter = 0\n",
    "\n",
    "        # Instead of list of tuples as the exp.replay concept go\n",
    "        # We use different np.arrays for each tuple element\n",
    "        self.state_buffer = np.zeros((self.buffer_capacity, num_states))\n",
    "        self.action_buffer = np.zeros((self.buffer_capacity, num_actions))\n",
    "        self.reward_buffer = np.zeros((self.buffer_capacity, 1))\n",
    "        self.next_state_buffer = np.zeros((self.buffer_capacity, num_states))\n",
    "\n",
    "    # Takes (s,a,r,s') obervation tuple as input\n",
    "    def record(self, obs_tuple):\n",
    "        # Set index to zero if buffer_capacity is exceeded,\n",
    "        # replacing old records\n",
    "        index = self.buffer_counter % self.buffer_capacity\n",
    "\n",
    "        self.state_buffer[index] = obs_tuple[0]\n",
    "        self.action_buffer[index] = obs_tuple[1]\n",
    "        self.reward_buffer[index] = obs_tuple[2]\n",
    "        self.next_state_buffer[index] = obs_tuple[3]\n",
    "\n",
    "        self.buffer_counter += 1\n",
    "\n",
    "    # We compute the loss and update parameters\n",
    "    def learn(self):\n",
    "        # Get sampling range\n",
    "        record_range = min(self.buffer_counter, self.buffer_capacity)\n",
    "        # Randomly sample indices\n",
    "        batch_indices = np.random.choice(record_range, self.batch_size)\n",
    "\n",
    "        # Convert to tensors\n",
    "        state_batch = tf.convert_to_tensor(self.state_buffer[batch_indices])\n",
    "        action_batch = tf.convert_to_tensor(self.action_buffer[batch_indices])\n",
    "        reward_batch = tf.convert_to_tensor(self.reward_buffer[batch_indices])\n",
    "        reward_batch = tf.cast(reward_batch, dtype=tf.float32)\n",
    "        next_state_batch = tf.convert_to_tensor(self.next_state_buffer[batch_indices])\n",
    "\n",
    "        # Training and updating Actor & Critic networks.\n",
    "        # See Pseudo Code.\n",
    "        with tf.GradientTape() as tape:\n",
    "            target_actions = target_actor(next_state_batch)\n",
    "            y = reward_batch + gamma * target_critic([next_state_batch, target_actions])\n",
    "            critic_value = critic_model([state_batch, action_batch])\n",
    "            critic_loss = tf.math.reduce_mean(tf.math.square(y - critic_value))\n",
    "\n",
    "        critic_grad = tape.gradient(critic_loss, critic_model.trainable_variables)\n",
    "        critic_optimizer.apply_gradients(\n",
    "            zip(critic_grad, critic_model.trainable_variables)\n",
    "        )\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            actions = actor_model(state_batch)\n",
    "            critic_value = critic_model([state_batch, actions])\n",
    "            # Used `-value` as we want to maximize the value given\n",
    "            # by the critic for our actions\n",
    "            actor_loss = -tf.math.reduce_mean(critic_value)\n",
    "\n",
    "        actor_grad = tape.gradient(actor_loss, actor_model.trainable_variables)\n",
    "        actor_optimizer.apply_gradients(\n",
    "            zip(actor_grad, actor_model.trainable_variables)\n",
    "        )\n",
    "\n",
    "\n",
    "# This update target parameters slowly\n",
    "# Based on rate `tau`, which is much less than one.\n",
    "def update_target(tau):\n",
    "    new_weights = []\n",
    "    target_variables = target_critic.weights\n",
    "    for i, variable in enumerate(critic_model.weights):\n",
    "        new_weights.append(variable * tau + target_variables[i] * (1 - tau))\n",
    "\n",
    "    target_critic.set_weights(new_weights)\n",
    "\n",
    "    new_weights = []\n",
    "    target_variables = target_actor.weights\n",
    "    for i, variable in enumerate(actor_model.weights):\n",
    "        new_weights.append(variable * tau + target_variables[i] * (1 - tau))\n",
    "\n",
    "    target_actor.set_weights(new_weights)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Here we define the Actor and Critic networks. These are basic Dense models\n",
    "with `ReLU` activation. `BatchNormalization` is used to normalize dimensions across\n",
    "samples in a mini-batch, as activations can vary a lot due to fluctuating values of input\n",
    "state and action.\n",
    "\n",
    "Note: We need the initialization for last layer of the Actor to be between\n",
    "`-0.003` and `0.003` as this prevents us from getting `1` or `-1` output values in\n",
    "the initial stages, which would squash our gradients to zero,\n",
    "as we use the `tanh` activation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "def get_actor():\n",
    "    # Initialize weights between -3e-3 and 3-e3\n",
    "    last_init = tf.random_uniform_initializer(minval=-0.003, maxval=0.003)\n",
    "\n",
    "    inputs = layers.Input(shape=(num_states,))\n",
    "    out = layers.Dense(512, activation=\"relu\")(inputs)\n",
    "    out = layers.BatchNormalization()(out)\n",
    "    out = layers.Dense(512, activation=\"relu\")(out)\n",
    "    out = layers.BatchNormalization()(out)\n",
    "    outputs = layers.Dense(1, activation=\"tanh\", kernel_initializer=last_init)(out)\n",
    "\n",
    "    # Our upper bound is 2.0 for Pendulum.\n",
    "    outputs = outputs * upper_bound\n",
    "    model = tf.keras.Model(inputs, outputs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def get_critic():\n",
    "    # State as input\n",
    "    state_input = layers.Input(shape=(num_states))\n",
    "    state_out = layers.Dense(16, activation=\"relu\")(state_input)\n",
    "    state_out = layers.BatchNormalization()(state_out)\n",
    "    state_out = layers.Dense(32, activation=\"relu\")(state_out)\n",
    "    state_out = layers.BatchNormalization()(state_out)\n",
    "\n",
    "    # Action as input\n",
    "    action_input = layers.Input(shape=(num_actions))\n",
    "    action_out = layers.Dense(32, activation=\"relu\")(action_input)\n",
    "    action_out = layers.BatchNormalization()(action_out)\n",
    "\n",
    "    # Both are passed through seperate layer before concatenating\n",
    "    concat = layers.Concatenate()([state_out, action_out])\n",
    "\n",
    "    out = layers.Dense(512, activation=\"relu\")(concat)\n",
    "    out = layers.BatchNormalization()(out)\n",
    "    out = layers.Dense(512, activation=\"relu\")(out)\n",
    "    out = layers.BatchNormalization()(out)\n",
    "    outputs = layers.Dense(1)(out)\n",
    "\n",
    "    # Outputs single value for give state-action\n",
    "    model = tf.keras.Model([state_input, action_input], outputs)\n",
    "\n",
    "    return model\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "`policy()` returns an action sampled from our Actor network plus some noise for\n",
    "exploration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "\n",
    "def policy(state, noise_object):\n",
    "    sampled_actions = tf.squeeze(actor_model(state))\n",
    "    noise = noise_object()\n",
    "    # Adding noise to action\n",
    "    sampled_actions = sampled_actions.numpy() + noise\n",
    "\n",
    "    # We make sure action is within bounds\n",
    "    legal_action = np.clip(sampled_actions, lower_bound, upper_bound)\n",
    "\n",
    "    return [np.squeeze(legal_action)]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Training hyperparameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "std_dev = 0.2\n",
    "ou_noise = OUActionNoise(mean=np.zeros(1), std_deviation=float(std_dev) * np.ones(1))\n",
    "\n",
    "actor_model = get_actor()\n",
    "critic_model = get_critic()\n",
    "\n",
    "target_actor = get_actor()\n",
    "target_critic = get_critic()\n",
    "\n",
    "# Making the weights equal initially\n",
    "target_actor.set_weights(actor_model.get_weights())\n",
    "target_critic.set_weights(critic_model.get_weights())\n",
    "\n",
    "# Learning rate for actor-critic models\n",
    "critic_lr = 0.002\n",
    "actor_lr = 0.001\n",
    "\n",
    "critic_optimizer = tf.keras.optimizers.Adam(critic_lr)\n",
    "actor_optimizer = tf.keras.optimizers.Adam(actor_lr)\n",
    "\n",
    "total_episodes = 100\n",
    "# Discount factor for future rewards\n",
    "gamma = 0.99\n",
    "# Used to update target networks\n",
    "tau = 0.005\n",
    "\n",
    "buffer = Buffer(50000, 64)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Now we implement our main training loop, and iterate over episodes.\n",
    "We sample actions using `policy()` and train with `learn()` at each time step,\n",
    "along with updating the Target networks at a rate `tau`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# To store reward history of each episode\n",
    "ep_reward_list = []\n",
    "# To store average reward history of last few episodes\n",
    "avg_reward_list = []\n",
    "\n",
    "# Takes about 20 min to train\n",
    "for ep in range(total_episodes):\n",
    "\n",
    "    prev_state = env.reset()\n",
    "    episodic_reward = 0\n",
    "\n",
    "    while True:\n",
    "        # Uncomment this to see the Actor in action\n",
    "        # But not in a python notebook.\n",
    "        # env.render()\n",
    "\n",
    "        tf_prev_state = tf.expand_dims(tf.convert_to_tensor(prev_state), 0)\n",
    "\n",
    "        action = policy(tf_prev_state, ou_noise)\n",
    "        # Recieve state and reward from environment.\n",
    "        state, reward, done, info = env.step(action)\n",
    "\n",
    "        buffer.record((prev_state, action, reward, state))\n",
    "        episodic_reward += reward\n",
    "\n",
    "        buffer.learn()\n",
    "        update_target(tau)\n",
    "\n",
    "        # End this episode when `done` is True\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "        prev_state = state\n",
    "\n",
    "    ep_reward_list.append(episodic_reward)\n",
    "\n",
    "    # Mean of last 40 episodes\n",
    "    avg_reward = np.mean(ep_reward_list[-40:])\n",
    "    print(\"Episode * {} * Avg Reward is ==> {}\".format(ep, avg_reward))\n",
    "    avg_reward_list.append(avg_reward)\n",
    "\n",
    "# Plotting graph\n",
    "# Episodes versus Avg. Rewards\n",
    "plt.plot(avg_reward_list)\n",
    "plt.xlabel(\"Episode\")\n",
    "plt.ylabel(\"Avg. Epsiodic Reward\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "![Graph](https://i.imgur.com/sqEtM6M.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "If training proceeds correctly, the average episodic reward will increase with time.\n",
    "\n",
    "Feel free to try different learning rates, `tau` values, and architectures for the\n",
    "Actor and Critic networks.\n",
    "\n",
    "The Inverted Pendulum problem has low complexity, but DDPG work great on many other\n",
    "problems.\n",
    "\n",
    "Another great environment to try this on is `LunarLandingContinuous-v2`, but it will take\n",
    "more episodes to obtain good results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# Save the weights\n",
    "actor_model.save_weights(\"pendulum_actor.h5\")\n",
    "critic_model.save_weights(\"pendulum_critic.h5\")\n",
    "\n",
    "target_actor.save_weights(\"pendulum_target_actor.h5\")\n",
    "target_critic.save_weights(\"pendulum_target_critic.h5\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Before Training:\n",
    "\n",
    "![before_img](https://i.imgur.com/ox6b9rC.gif)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "After 100 episodes:\n",
    "\n",
    "![after_img](https://i.imgur.com/eEH8Cz6.gif)\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "ddpg_pendulum",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}